import hashlib
import os
import sqlite3
import sys

from loguru import logger


def calculate_md5(file_path):
    """Calculates the MD5 hash of a file."""
    try:
        md5 = hashlib.md5()
        with open(file_path, "rb") as f:
            md5.update(f.read())
        return md5.hexdigest()
    except Exception as e:
        print(f"Error calculating MD5 for {file_path}: {e}")
        return None


def migrate(root_dir):
    """Finds files, calculates MD5, and updates local_url in the database."""
    db_path = os.path.join(root_dir, "Cache.dbst3")

    if not os.path.exists(db_path):
        logger.error(f"Error: Database file not found at {db_path}")
        return

    conn = None
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        logger.debug(f"Successfully connected to SQLite database at {db_path}")

        file_types = (".md", ".png", ".jpg", ".jpeg", ".gif")
        updated_count = 0

        for subdir, _, files in os.walk(root_dir):
            for file in files:
                file_path = os.path.join(subdir, file)
                if file_path.lower().endswith(file_types):
                    relative_path = os.path.relpath(file_path, root_dir)
                    md5_hash = calculate_md5(file_path)

                    if md5_hash:
                        # Check if md5 exists in the cache table
                        cursor.execute(
                            "SELECT md5 FROM cache WHERE md5 = ?", (md5_hash,)
                        )
                        row = cursor.fetchone()

                        if row:
                            # If md5 exists, update the local_url
                            try:
                                cursor.execute(
                                    "UPDATE cache SET local_url = ? WHERE md5 = ?",
                                    (relative_path, md5_hash),
                                )
                                conn.commit()
                                updated_count += 1
                                logger.debug(
                                    f"Updated local_url for md5 {md5_hash} with {relative_path}"
                                )
                            except sqlite3.Error as e:
                                logger.debug(
                                    f"Error updating local_url for md5 {md5_hash}: {e}"
                                )
                        # else: print(f"MD5 {md5_hash} not found in database. Skipping {relative_path}") # Optional: log skipped files

        logger.info(
            f"Finished updating local_urls. Total updated entries: {updated_count}"
        )

    except sqlite3.Error as e:
        logger.error(f"SQLite database error: {e}")
        if conn:
            conn.rollback()
            logger.error("Transaction rolled back due to error.")
    except Exception as e:
        logger.error(f"An unexpected error occurred: {e}")
    finally:
        if conn:
            conn.close()
            logger.debug("SQLite database connection closed.")


if __name__ == "__main__":
    # Check if root directory path is provided as a command-line argument
    if len(sys.argv) < 2:
        logger.info(
            "Usage: python 002._add_comumn_local_url_in_table_cache.py <directory_path>"
        )
        sys.exit(1)
    directory_path = sys.argv[1]
    migrate(directory_path)
